from fastmcp import FastMCP, Client
from serpapi import GoogleSearch
from typing import Annotated
import json
import logging
import asyncio
import os
os.environ["SERP_KEY"] = ""
api_key = os.getenv("SERP_KEY")

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

mcp = FastMCP('demo.mcp')

@mcp.tool()
async def get_news_data(
    query: Annotated[str, "Search query, typically a keyword or phrase"],
    gl: Annotated[str, "Country code for localization (e.g., 'cn', 'us', 'uk')"] = "cn",
    hl: Annotated[str, "Language code for localization (e.g., 'zh-cn', 'en', 'es')"] = "zh-cn",
    num: Annotated[int, "Number of results to return (default: 10)"] = 10
) -> str:
    """
    Retrieves Google News results for a given query using the SerpApi Google News API.

    Input Parameters:
        query: The search query (e.g., Artificial Intelligence).
        gl: Country code for localization (default: 'cn' for China).
        hl: Language code for localization (default: 'zh-cn' for Chinese).
        num: Number of results to return (default: 10).

    Output:
        result: A JSON string containing a list of news articles with fields: link, title, snippet, date, and source.
                Returns an error message if the API call fails or no data is found.
    """
    # Validate inputs
    if not query or not query.strip():
        return json.dumps({"error": "Query cannot be empty."})
    if not gl or not hl:
        return json.dumps({"error": "Country code (gl) and language code (hl) cannot be empty."})
    if num <= 0:
        return json.dumps({"error": "Number of results (num) must be positive."})

    # SerpApi parameters
    params = {
        "engine": "google_news",
        "q": query,
        "api_key": api_key, 
        "hl": hl,
        "gl": gl,
        "num": num
    }

    news_results = []
    try:
        # Make API call using serpapi client
        logger.info(f"Making SerpApi request for query: {query}")
        search = GoogleSearch(params)
        results = search.get_dict()

        # Check search status
        if results.get("search_metadata", {}).get("status") != "Success":
            error_msg = results.get("error", "Unknown API error")
            logger.error(f"API error: {error_msg}")
            return json.dumps({"error": f"API error: {error_msg}"})

        # Extract news results
        for item in results.get("news_results", []):
            try:
                news_results.append({
                    "link": item.get("link", ""),
                    "title": item.get("title", ""),
                    "snippet": item.get("snippet", ""),
                    "date": item.get("date", ""),
                    "source": item.get("source", {}).get("name", "")
                })
            except Exception as e:
                logger.warning(f"Error processing news result: {e}")
                continue

    except Exception as e:
        logger.error(f"Unexpected error during API call: {e}")
        return json.dumps({"error": f"Unexpected error during API call: {str(e)}"})

    # Return results as JSON string
    if not news_results:
        return json.dumps({"error": f"No news results found for query '{query}'"})
    return json.dumps(news_results, indent=2, ensure_ascii=False)  # ensure_ascii=False for proper Chinese character encoding

GOOGLE_NEWS_SCHEMA = {
    "type": "function",
    "function": {
        "name": "get_news_data",
        "description": "Retrieves Google News results for a given query using the SerpApi Google News API.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search query (e.g., Artificial Intelligence)."
                },
                "gl": {
                    "type": "string",
                    "description": "Country code for localization (e.g., 'cn', 'us', 'uk').",
                    "default": "cn"
                },
                "hl": {
                    "type": "string",
                    "description": "Language code for localization (e.g., 'zh-cn', 'en', 'es').",
                    "default": "zh-cn"
                },
                "num": {
                    "type": "integer",
                    "description": "Number of results to return.",
                    "default": 10
                }
            },
            "required": ["query"]
        }
    }
}

async def main():
    """
    Main entry point for the FastMCP client to interact with the Google News API tool.
    Initializes the client, lists available tools, and calls the news data tool with sample parameters.
    """
    client = Client(mcp)
    async with client:
        # List available tools for debugging/verification
        tools = await client.list_tools()
        print('Available tools:', tools)

        # Call the news data tool with sample parameters
        result = await client.call_tool('get_news_data', {
            'query': 'Artificial Intelligence',
            'gl': 'cn',
            'hl': 'zh-cn',
            'num': 10
        })
        print('News Data Result:', result)

# Recommended entry point for running the async program
if __name__ == '__main__':
    asyncio.run(main())